/**
 * Copyright 2019-2022 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <gtest/gtest.h>
#include <mockcpp/mockcpp.hpp>
#include <mockcpp/mockable.h>
#include <dlfcn.h>
#include "graph/graph.h"
#include "graph/model.h"
#include "graph/operator_hiai_reg.h"
#include "graph/op/all_ops.h"
#include "graph/op/image_para_utils.h"
#include "framework/tensor/image_format.h"
#include "framework/tensor/image_process_config_types.h"
#include "graph/op/aipp_defs.h"

using testing::Test;
using namespace std;
using namespace hiai;
using namespace ge;

class GeImageParaUtilUnittest : public testing::Test {
public:
    void SetUp()
    {
    }

    void TearDown()
    {
    }
};

hiai::op::Data CreateData(Shape& shape, const std::string& name)
{
    ge::TensorDesc dataDesc(shape, FORMAT_NCHW, DT_UINT8);
    hiai::op::Data data = hiai::op::Data(name);
    data.update_input_desc_x(dataDesc);
    return data;
}

/*
 * 测试用例名称   : image_para_util_aipp_dynamic_func_all
 * 测试用例描述： ir api 构建graph, 所有7个动态AIPP算子
 * 预置条件 :ir api 构建模型
 * 预期结果 :构建模型成功
 */
TEST_F(GeImageParaUtilUnittest, image_para_util_aipp_dynamic_func_all)
{
    ge::Graph graph("graph_defalut");
    Shape shape({1, 4, 200, 320});
    hiai::op::Data data = CreateData(shape, "data1");
    hiai::MultiCropPara multiCrop;
    multiCrop.cropParas.resize(1);

    hiai::op::ConfigData cropConfigData = ImageParaUtil::CreateConfigData(
        multiCrop, "cropConfigData", hiai::op::ImageMultiCrop::TYPE);
    hiai::op::Data data10 = CreateData(shape, "data10");
    cropConfigData.set_input_x(data10);

    auto imageMultiCrop = hiai::op::ImageMultiCrop("imageMultiCrop")
                       .set_input_x(data)
                       .set_input_param(cropConfigData);

    ChannelSwapPara channelswp_ara;
    channelswp_ara.rbuvSwapSwitch = true;
    channelswp_ara.axSwapSwitch = true;
    hiai::op::ConfigData channelSwapConfigData = ImageParaUtil::CreateConfigData(
        channelswp_ara, "channelSwapConfigData", hiai::op::ImageChannelSwapV2::TYPE);
    hiai::op::Data data11 = CreateData(shape, "data11");
    channelSwapConfigData.set_input_x(data11);

    auto channelSwapv2 = hiai::op::ImageChannelSwapV2("channelSwapv2")
                       .set_input_x(imageMultiCrop)
                       .set_input_param(channelSwapConfigData);

    CscPara csc_para;
    csc_para.outputFormat = ImageFormat::YUV444SP;
    csc_para.imageColorSpace = ImageColorSpace::JPEG;
    hiai::op::ConfigData cscConfigData = ImageParaUtil::CreateConfigData(
        csc_para, "cscConfigData", hiai::op::ImageColorSpaceConvertionV2::TYPE);
    hiai::op::Data data12 = CreateData(shape, "data12");
    cscConfigData.set_input_x(data12);

    auto cscV2 = hiai::op::ImageColorSpaceConvertionV2("cscV2")
                       .set_input_x(channelSwapv2)
                       .set_input_param(cscConfigData);

    MultiResizePara multiResizePara;
    multiResizePara.resizeParas.resize(1);
    hiai::op::ConfigData resizeConfigData = ImageParaUtil::CreateConfigData(
        multiResizePara, "resizeConfigData", hiai::op::ImageMultiResize::TYPE);
    hiai::op::Data data13 = CreateData(shape, "data13");
    resizeConfigData.set_input_x(data13);

    auto imageMultiResize = hiai::op::ImageMultiResize("imageMultiResize")
                       .set_input_x(cscV2)
                       .set_input_param(resizeConfigData);

    MultiDtcPara multiDtcPara;
    multiDtcPara.dtcParas.resize(1);

    hiai::op::ConfigData dtcConfigData = ImageParaUtil::CreateConfigData(
        multiDtcPara, "dtcConfigData", hiai::op::ImageMultiDataTypeConvertion::TYPE);
    hiai::op::Data data14 = CreateData(shape, "data14");
    dtcConfigData.set_input_x(data14);

    auto imageMultiDataTypeConvertion = hiai::op::ImageMultiDataTypeConvertion("imageMultiDataTypeConvertion")
                       .set_input_x(imageMultiResize)
                       .set_input_param(dtcConfigData);

    MultiRotatePara multiRotatePara;
    multiRotatePara.rotateParas.resize(1);

    hiai::op::ConfigData rotateConfigData = ImageParaUtil::CreateConfigData(
        multiRotatePara, "rotateConfigData", hiai::op::ImageMultiRotate::TYPE);
    hiai::op::Data data15 = CreateData(shape, "data15");
    rotateConfigData.set_input_x(data15);

    auto imageMultiRotatePara = hiai::op::ImageMultiRotate("imageMultiRotatePara")
                       .set_input_x(imageMultiDataTypeConvertion)
                       .set_input_param(rotateConfigData);

    MultiPadPara multiPadPara;
    multiPadPara.padParas.resize(1);
    hiai::op::ConfigData padConfigData = ImageParaUtil::CreateConfigData(
        multiPadPara, "padConfigData", hiai::op::ImageMultiPad::TYPE);
    hiai::op::Data data16 = CreateData(shape, "data16");
    padConfigData.set_input_x(data16);

    auto imageMultiPad = hiai::op::ImageMultiPad("imageMultiPad")
                       .set_input_x(imageMultiRotatePara)
                       .set_input_param(padConfigData);

    std::vector<Operator> inputs{data, cropConfigData, channelSwapConfigData,
        cscConfigData, resizeConfigData, dtcConfigData, rotateConfigData, padConfigData};
    std::vector<Operator> outputs{imageMultiPad};
    graph.SetInputs(inputs);
    graph.SetOutputs(outputs);

    EXPECT_EQ(true, graph.IsValid());
}


/*
 * 测试用例名称   : image_para_util_aipp_static_func_all
 * 测试用例描述： ir api 构建graph, 7个静态aipp算子
 * 预置条件 :ir api 构建模型
 * 预期结果 :构建模型成功
 */
TEST_F(GeImageParaUtilUnittest, image_para_util_aipp_static_func_all)
{
    ge::Graph graph("graph_defalut");
    Shape shape({1, 4, 200, 320});
    hiai::op::Data data = CreateData(shape, "data1");
    CropPara cropPara;
    cropPara.cropStartPosW = 0;
    cropPara.cropStartPosH = 0;
    cropPara.cropSizeW = 240;
    cropPara.cropSizeH = 160;

    hiai::MultiCropPara multiCrop;
    multiCrop.imageFormat = ImageFormat::XRGB8888;
    multiCrop.cropParas.push_back(cropPara);

    hiai::op::Const cropConstData = ImageParaUtil::CreateConfigConst(multiCrop, "cropConstData");
    auto imageMultiCrop = hiai::op::ImageMultiCrop("imageMultiCrop")
                       .set_input_x(data)
                       .set_input_param(cropConstData);

    ChannelSwapPara channelswp_ara;
    channelswp_ara.imageFormat = ImageFormat::RGB888;
    channelswp_ara.rbuvSwapSwitch = true;
    channelswp_ara.axSwapSwitch = true;
    hiai::op::Const channelSwapConstData = ImageParaUtil::CreateConfigConst(channelswp_ara, "channelSwapConstData");
    auto channelSwapv2 = hiai::op::ImageChannelSwapV2("channelSwapv2")
                       .set_input_x(imageMultiCrop)
                       .set_input_param(channelSwapConstData);

    CscPara csc_para;
    csc_para.outputFormat = ImageFormat::YUV444SP;
    csc_para.imageColorSpace = ImageColorSpace::JPEG;
    hiai::op::Const cscConstData = ImageParaUtil::CreateConfigConst(csc_para, "cscConstData");
    auto cscV2 = hiai::op::ImageColorSpaceConvertionV2("cscV2")
                       .set_input_x(channelSwapv2)
                       .set_input_param(cscConstData);

    ResizePara resizePara;
    resizePara.resizeOutputSizeW = 320;
    resizePara.resizeOutputSizeH = 150;
    MultiResizePara multiResizePara;
    multiResizePara.resizeParas.push_back(resizePara);
    hiai::op::Const resizeConstData = ImageParaUtil::CreateConfigConst(multiResizePara, "resizeConstData");
    auto imageMultiResize = hiai::op::ImageMultiResize("imageMultiResize")
                       .set_input_x(cscV2)
                       .set_input_param(resizeConstData);

    MultiDtcPara multiDtcPara;
    multiDtcPara.dtcParas.resize(1);

    hiai::op::Const dtcConstData = ImageParaUtil::CreateConfigConst(multiDtcPara, "dtcConstData");
    auto imageMultiDataTypeConvertion = hiai::op::ImageMultiDataTypeConvertion("imageMultiDataTypeConvertion")
                       .set_input_x(imageMultiResize)
                       .set_input_param(dtcConstData);

    RotatePara rotatePara;
    rotatePara.rotationAngle = 90.0f;

    MultiRotatePara multiRotatePara;
    multiRotatePara.rotateParas.push_back(rotatePara);

    hiai::op::Const rotateConstData = ImageParaUtil::CreateConfigConst(multiRotatePara, "rotateConstData");
    auto imageMultiRotate = hiai::op::ImageMultiRotate("imageMultiRotate")
                       .set_input_x(imageMultiDataTypeConvertion)
                       .set_input_param(rotateConstData);

    PadPara padding_para;
    padding_para.paddingSizeTop = 10;
    padding_para.paddingSizeBottom = 10;
    padding_para.paddingSizeLeft = 20;
    padding_para.paddingSizeRight = 20;

    MultiPadPara multiPadPara;
    multiPadPara.padParas.push_back(padding_para);

    hiai::op::Const padConstData = ImageParaUtil::CreateConfigConst(multiPadPara, "padConstData");
    auto imageMultiPad = hiai::op::ImageMultiPad("imageMultiPad")
                       .set_input_x(imageMultiRotate)
                       .set_input_param(padConstData);

    std::vector<Operator> inputs{data};
    std::vector<Operator> outputs{imageMultiPad};
    graph.SetInputs(inputs);
    graph.SetOutputs(outputs);

    EXPECT_EQ(true, graph.IsValid());
}
